#!/usr/bin/env python
import re
import StringIO
import struct
import ply.lex as lex
import ply.yacc as yacc

class Assembler:
    # List of token names.
    tokens = (
       'LABEL',
       'MNEMONIC',
       'NUMBER',
       'REGISTER',
       'COMMENT',
       'COMMA',
    )

    # List of mnemonics and their opcodes
    mnemonic = {
        #r-type instructions
        'ADD':  0b00010001,
        'AND':  0b00010010,
        'LDR':  0b00010011,

        #i-type instructions
        'ADDI': 0b00100001,
        'ANDI': 0b00100010,
        'LDRI': 0b00100011,
    }

    # Ignore whitespace and comments
    t_ignore  = ' \t' 
    t_ignore_COMMENT = '\#.*'
    t_COMMA  = ',' 

    def __init__(self):
        self.lineno = 0
        # Build the lexer
        self.lexer = lex.lex(module=self)
        # Build the parser
        self.parser = yacc.yacc(module=self)
    
    def lexical(self, l):
        # Give the lexer some input
        self.lexer.input(data)

        # Tokenize
        while True:
            tok = self.lexer.token()
            if not tok: break      # No more input
            print tok.type, tok.value, tok.lineno

    def asm(self, data):
        buffer = StringIO.StringIO()

        for l in re.split('(\n)', data):
            result = self.parser.parse(l)
            if result == None:
                self.p_error(None)
                buffer.close()
                return None
            buffer.write(result) 

        result = buffer.getvalue()
        buffer.close()            
        return result

    # Rule for tracking line numbers
    def t_newline(self, t):
        r'\n+'
        self.lexer.lineno += len(t.value)

    # Rule for error handling
    def t_error(self, t):
        print "Illegal character '%s'" % t.value[0]
        t.lexer.skip(1)

    # Rule for matching labels
    def t_LABEL(self, t):
        r'[a-zA-Z_0-9]+:'
        t.value = t.value[:-1]
        return t

    # Rule for matching numbers
    def t_NUMBER(self, t):
        r'0x[a-fA-F0-9]+|[0-9]+'
        t.value = int(t.value, 0)
        return t

    # Rule for matching registers
    def t_REGISTER(self, t):
        r'r[0-9]'
        t.value = int(t.value[1], 0)
        return t

    # Rule for matching mnemonics
    def t_MNEMONIC(self, t):
        r'[a-zA-Z]+'
        t.value = self.mnemonic[t.value]
        return t

    # parser rules
    def p_instruction(self, p):
        '''instruction : r_inst
                       | i_inst
                       | empty
        '''
        p[0] = p[1]

    def p_empty(self, p):
        '''empty :'''
        p[0] = ''
        
    def p_r_instruction(self, p):
        'r_inst : MNEMONIC REGISTER COMMA REGISTER'
        if p[1] & 0x10:
            p[0] = struct.pack('4B', p[1],p[2],p[4], 0)

    def p_i_instruction(self, p):
        'i_inst : MNEMONIC REGISTER COMMA NUMBER'
        if p[4] > 0xFFFF:
            print 'Immediate operand out of range %d' %self.lineno
            return         
        if p[1] & 0x20:
            p[0] = struct.pack('2BH', p[1],p[2],p[4])

    def p_operand(self, p):
        '''operand : NUMBER 
                   | REGISTER'''
        print p[1]

    def p_error(self, p):
        print 'Error parsing input! lineno',self.lexer.lineno

if __name__ == "__main__":
    data = ''' 
        LDRI r0, 16000
        LDRI r1, 0xDEAD
        LDRI r2, 0xBEEF
        ADDI r0, 1
        # comment

    '''
    asm = Assembler()
    with open('output.bin', 'wb') as f:
        res = asm.asm(data)
        if res:
            f.write(res)

